package cn.akwangl.socket.nio;

import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Set;

/**
 * Socket NIO 模型测试.Server 端实现.
 *
 * @author 阿坤同学
 * @date 2022/6/22 14:05
 * @since 0.0.6
 */
@Slf4j
public class NIOChatService {
    public static final int DEFAULT_BUFFER_SIZE = 1024;
    public static final int DEFAULT_SERVER_PORT = 7000;

    public static final String QUIT = "quit";

    public static final String SERVER_ADDRESS = "localhost";
    private Charset charset = StandardCharsets.UTF_8;


    private ServerSocketChannel serverSocket;
    private Selector selector;
    private ByteBuffer readBuffer = ByteBuffer.allocate(DEFAULT_BUFFER_SIZE);
    private ByteBuffer writeBuffer = ByteBuffer.allocate(DEFAULT_BUFFER_SIZE);

    public void start() {
        try {
            serverSocket = ServerSocketChannel.open();
            serverSocket.bind(new InetSocketAddress(SERVER_ADDRESS, DEFAULT_SERVER_PORT));
            serverSocket.configureBlocking(false);

            selector = Selector.open();
            serverSocket.register(selector, SelectionKey.OP_ACCEPT);

            while (true) {
                selector.select();
                Set<SelectionKey> selectionKeys = selector.selectedKeys();

                for (SelectionKey selectionKey : selectionKeys) {
                    eventHandlers(selectionKey);
                }
                selectionKeys.clear();
            }
        } catch (IOException e) {
            log.error("server socket 连接失败.");
        } finally {
            close();
        }
    }

    /**
     * 事件处理器.
     * <p>
     * 当一个客户端连接服务器时首先会监听{@link SelectionKey#OP_ACCEPT}事件,{@link SelectionKey#OP_ACCEPT}事件处理时会将
     * 客户端Channel注册为{@link SelectionKey#OP_READ},当客户端有新的数据发往服务器,服务器群发给当前所有在线用户.
     *
     * @param selectionKey 选择器
     * @author 阿坤同学
     * @date 2022/6/23 14:25
     * @since 0.0.6
     */
    private void eventHandlers(SelectionKey selectionKey) throws IOException {
        // 处理客户端连接请求
        if (selectionKey.isAcceptable()) {
            ServerSocketChannel serverSocket = (ServerSocketChannel) selectionKey.channel();
            SocketChannel clientChannel = serverSocket.accept();
            clientChannel.configureBlocking(false);

            // 注册读取事件
            clientChannel.register(selector, SelectionKey.OP_READ);
            log.info("service端:[{}] 已连接", clientChannel.socket().getRemoteSocketAddress());
        } else if (selectionKey.isReadable()) {
            SocketChannel clientChannel = (SocketChannel) selectionKey.channel();
            SocketAddress remoteSocketAddress = clientChannel.socket().getRemoteSocketAddress();

            String clientMessage = receive(clientChannel);
            clientChannel.register(selector, SelectionKey.OP_WRITE);
            if (clientMessage != null && clientMessage.length() > 0) {
                log.info("service 接收到的消息:{}", clientMessage);

                // 发送消息到其他客户端
                forwardMessage(clientChannel, clientMessage);

                // 退出
                if (QUIT.equals(clientMessage)) {
                    selectionKey.cancel();
                    selector.wakeup();
                    log.info("service 客户端:{} 已断开连接", remoteSocketAddress);
                }
            } else {
                selectionKey.channel();
                selector.wakeup();
            }
        }
    }

    /**
     * 群发.将消息转发给其他在线客户端.
     *
     * @param clientChannel 排除接收的客户端
     * @param clientMessage 客户端消息
     * @author 阿坤同学
     * @date 2022/6/23 22:12
     * @since 0.0.6
     */
    private void forwardMessage(SocketChannel clientChannel, String clientMessage) throws IOException {
        Set<SelectionKey> selectionKeys = selector.keys();
        for (SelectionKey selectionKey : selectionKeys) {
            if (selectionKey.channel() instanceof ServerSocketChannel) {
                continue;
            }
            if (selectionKey.isValid() && !clientChannel.equals(selectionKey.channel())) {
                writeBuffer.clear();
                writeBuffer.put(charset.encode(clientMessage));
                writeBuffer.flip();
                // 发送数据给其他客户端.
                while (writeBuffer.hasRemaining()) {
                    ((SocketChannel) selectionKey.channel()).write(writeBuffer);
                }
            }
        }
    }

    /**
     * 接收客户端的消息.
     *
     * @param clientChannel 客户端Channel
     * @return 客户端数据
     * @author 阿坤同学
     * @date 2022/6/23 22:12
     * @since 0.0.6
     */
    private String receive(SocketChannel clientChannel) throws IOException {
        readBuffer.clear();
        while (clientChannel.read(readBuffer) > 0) ;
        readBuffer.flip();
        return String.valueOf(charset.decode(readBuffer));
    }

    public void close() {
        if (!serverSocket.isOpen()) {
            try {
                serverSocket.close();
                log.info("server socket 已断开连接...");
            } catch (IOException e) {
                log.error("socket 连接关闭失败", e);
            }
        }
    }

    public static void main(String[] args) {
        NIOChatService chatService = new NIOChatService();
        chatService.start();
    }
}
